#include <stdlib.h>
#include <stdio.h>
#include <stdarg.h>
#include <assert.h>
#include <math.h>
#include "tabular.h"

void TabularCPD_Init_Params(TabularCPD_t * dist, int nparams, ...)
{
  int i;
  va_list pars;
  
  dist->numrows = nparams;
  dist->p_rows = (RkVector_t **) malloc(sizeof(*dist->p_rows) * nparams);
  
  va_start(pars, nparams);

  for (i=0; i<nparams; i++)
  {
    dist->p_rows[i] = va_arg(pars, RkVector_t *);
    if (0 == i)
      dist->numcols = RkVector_SIZE(dist->p_rows[i]);
    else if (dist->numcols != RkVector_SIZE(dist->p_rows[i]))
    {
      fprintf(stderr, "error: TabularCPD requires RkVectors of same size\n");
      exit(1);
    }

    if (fabs(1-RkVector_sum(dist->p_rows[i])) > 1e-4)
    {
      fprintf(stderr, "error: TabularCPD RkVector num %d doesn't sum to 1"
              " (sums to %f)\n", i+1, RkVector_sum(dist->p_rows[i]));
      exit(1);
    }
  }

  va_end(pars);
}

/* give the cardinality of each argument */
void TabularCPD_Init_Args(TabularCPD_t * dist, int nargs, ...)
{
  /* Init_Params should have been called before Init_Args */
  assert(dist->numrows && dist->p_rows && dist->numcols);

  dist->numargs = nargs;

  if (dist->numargs > 0)
  {
    int i;
    int numrows;
    va_list args;

    dist->p_arg_card = (int *) malloc(sizeof(*dist->p_arg_card) * nargs);

    va_start(args, nargs);
    
    for (i=0; i<dist->numargs; i++)
      dist->p_arg_card[i] = va_arg(args, int);

    va_end(args);

    /* check that the number of rows is correct */
    numrows = 1;
    for (i=0; i<dist->numargs; i++)
    {
      if (dist->p_arg_card[i] <= 0)
      {
        fprintf(stderr, "error: TabularCPD requires all arguments of +ve "
                "cardinality");
        exit(1);
      }
      numrows *= dist->p_arg_card[i];
    }
    
    if (numrows != dist->numrows)
    {
      fprintf(stderr, "error: TabularCPD requires one RkVector for each "
              "combination of argument values");
      exit(1);
    }
  }
  else if (0 == dist->numargs)
  {
    dist->p_arg_card = NULL;
  }
  else
  {
    fprintf(stderr, "error: TabularCPD number of arguments must be >= 0");
    exit(1);
  }
}

static int TabularCPD_rownum(const TabularCPD_t * dist, int nargs,
                             va_list args)
{
  assert(dist->numargs == nargs);

  if (!nargs)
  {
    return 0;
  }
  else
  {
    int i;
    int idx;
    
    idx = va_arg(args, int);

    for (i=1; i<nargs; i++)
    {
      idx *= dist->p_arg_card[i];
      idx += va_arg(args, int);
    }
    
    return idx;
  }
}

int TabularCPD_Sample(const TabularCPD_t * dist, int nargs, ...)
{
  va_list args;
  int idx;
  
  if (!dist->numcols)
    return -1;
  
  va_start(args, nargs);
  idx = TabularCPD_rownum(dist, nargs, args);
  va_end(args);

  return RkVector_sample(dist->p_rows[idx]);
}

double TabularCPD_Prob(const TabularCPD_t * dist, int val, int nargs, ...)
{
  va_list args;
  int idx;
  
  assert((val >= -1) && (val < dist->numcols));

  if (!dist->numcols && (-1 == val))
    return 1;
  else if (-1 == val)
    return 0;
  
  va_start(args, nargs);
  idx = TabularCPD_rownum(dist, nargs, args);
  va_end(args);
  
  return RkVector_ELEM(dist->p_rows[idx], val);
}

double TabularCPD_LogProb(const TabularCPD_t * dist, int val, int nargs, ...)
{
  va_list args;
  int idx;
  
  assert((val >= -1) && (val < dist->numcols));

  if (!dist->numcols && (-1 == val))
    return 0;
  else if (-1 == val)
    return -HUGE_VAL;
  
  va_start(args, nargs);
  idx = TabularCPD_rownum(dist, nargs, args);
  va_end(args);

  return log(RkVector_ELEM(dist->p_rows[idx], val));
}
